#include <rclcpp/rclcpp.hpp>
#include <tf2_ros/transform_listener.h>
#include <tf2_ros/transform_broadcaster.h>
#include <geometry_msgs/msg/transform_stamped.hpp>
#include <tf2_ros/buffer.h>
#include <geometry_msgs/msg/twist.hpp>
#include <mutex>

/*
此代码作用：读取相机和tag、机械臂基坐标系和机械臂末端的TF变换，发布重命名后的新的TF变换

相机原始TF变换：camera_color_optical_frame -> aruco_marker_frame
重命名后的TF变换：virtual -> tag

机械臂原始TF变换：world -> link_eef
重命名后的TF变换：virtual -> ee

小乌龟控制：按下键盘上的上键，fake_origin的TF变换不再更新，固定在当前位置，作为控制的参考零点
*/

using namespace std::chrono_literals;

class FrameListener : public rclcpp::Node
{
public:
    FrameListener()
    : Node("tf2_frame_listener")
    {
    tf_buffer_ =
      std::make_unique<tf2_ros::Buffer>(this->get_clock());
    tf_listener_ =
      std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);
    tf_broadcaster_ =
      std::make_unique<tf2_ros::TransformBroadcaster>(*this);
    subscription_ = this->create_subscription<geometry_msgs::msg::Twist>(
            "turtle1/cmd_vel", 10, std::bind(&FrameListener::teleopCallback, this, std::placeholders::_1));
    
    tf_publisher_thread_ = std::thread(&FrameListener::tfPublisherThread, this);

    tf_listener_thread_ = std::thread(&FrameListener::tfListenerThread, this);
    }

private:

void teleopCallback(const geometry_msgs::msg::Twist::SharedPtr msg)
{
    if(msg->linear.x > 0.1) flag_ = true;
    if(msg->linear.y > 0.1) flag_ = false;
}
    bool flag_=false;
    void tfPublisherThread();
    void tfListenerThread();

    rclcpp::TimerBase::SharedPtr timer_{nullptr};
    std::shared_ptr<tf2_ros::TransformListener> tf_listener_{nullptr};
    std::unique_ptr<tf2_ros::TransformBroadcaster> tf_broadcaster_;
    std::unique_ptr<tf2_ros::Buffer> tf_buffer_;
    rclcpp::Subscription<geometry_msgs::msg::Twist>::SharedPtr subscription_;
    rclcpp::Publisher<geometry_msgs::msg::Twist>::SharedPtr publisher_;

    //char flag;
    std::thread tf_publisher_thread_;
    std::thread tf_listener_thread_;

    geometry_msgs::msg::TransformStamped original_transform_camera_; 
    geometry_msgs::msg::TransformStamped original_transform_arm_;
    geometry_msgs::msg::TransformStamped transform_memory;
    geometry_msgs::msg::TransformStamped transform_fake2tag_;
    std::mutex tf_mutex_;

};

void FrameListener::tfPublisherThread()
{
    while (rclcpp::ok())
    {
        rclcpp::sleep_for(std::chrono::milliseconds(20));
        geometry_msgs::msg::TransformStamped original_transform_camera_temp;
        geometry_msgs::msg::TransformStamped original_transform_arm_temp;

        {
            std::lock_guard<std::mutex> lock(tf_mutex_);
            original_transform_camera_temp = original_transform_camera_;
            original_transform_arm_temp = original_transform_arm_;
        }

        original_transform_camera_temp.header.frame_id = "virtual";
        original_transform_camera_temp.child_frame_id = "tag";
        // 设置时间戳等必要信息
        // original_transform_camera_temp.header.stamp = this->get_clock()->now();
        // 发布新的TF变换
        tf_broadcaster_->sendTransform(original_transform_camera_temp);

        // std::cout <<"test 4" << std::endl;


        // 对获取到的变换重新命名，这里将 "source_frame" 重命名为 "new_source_frame"，
        // 将 "target_frame" 重命名为 "new_target_frame"，可按需修改
        original_transform_arm_temp.header.frame_id = "virtual";
        original_transform_arm_temp.child_frame_id = "ee";
        // 设置时间戳等必要信息
        // original_transform_arm_temp.header.stamp = this->get_clock()->now();
        // 发布新的TF变换
        tf_broadcaster_->sendTransform(original_transform_arm_temp);

        if (flag_ == false) // up is pressed, fake is not updated
        {
            transform_memory = original_transform_camera_temp;
            transform_memory.header.frame_id = "virtual";
            transform_memory.child_frame_id = "fake_origin";

        }
        //设置时间戳等必要信息
        transform_memory.header.stamp = this->get_clock()->now();
        //发布新的TF变换
        tf_broadcaster_->sendTransform(transform_memory);


    }
}

void FrameListener::tfListenerThread()
{
    while(rclcpp::ok())
    {
        rclcpp::sleep_for(std::chrono::milliseconds(50));
        try {
        {    
            std::lock_guard<std::mutex> lock(tf_mutex_);
            original_transform_camera_ = tf_buffer_->lookupTransform(
                "camera_color_optical_frame", "aruco_marker_frame", tf2::TimePointZero);
        }
        } catch (const tf2::TransformException & ex) {
            RCLCPP_INFO(
                this->get_logger(), "Could not transform %s to %s: %s",
                "camera_color_optical_frame", "aruco_marker_frame", ex.what());
                rclcpp::sleep_for(std::chrono::milliseconds(1));
                continue;
            }


        try {
        // 假设原TF变换是从 "source_frame" 到 "target_frame"，这里根据实际情况修改
                {
                    std::lock_guard<std::mutex> lock(tf_mutex_);
                    original_transform_arm_ = tf_buffer_->lookupTransform(
                        "world", "link_eef", tf2::TimePointZero);
                }
            } 
        catch (const tf2::TransformException & ex) 
            {
                RCLCPP_INFO(
                    this->get_logger(), "Could not transform %s to %s: %s",
                    "camera_color_optical_frame", "aruco_marker_frame", ex.what());
                    rclcpp::sleep_for(std::chrono::milliseconds(1));
                    continue;
            }

        try {
        // 假设原TF变换是从 "source_frame" 到 "target_frame"，这里根据实际情况修改
        {   std::lock_guard<std::mutex> lock(tf_mutex_);
            transform_fake2tag_ = tf_buffer_->lookupTransform(
                "fake_origin","tag", tf2::TimePointZero);
        }
        } catch (const tf2::TransformException & ex) {
            RCLCPP_INFO(
                this->get_logger(), "Could not transform %s to %s: %s",
                "camera_color_optical_frame", "aruco_marker_frame", ex.what());
                rclcpp::sleep_for(std::chrono::milliseconds(1));
                continue;
            // return;
            }
    }
    
}

int main(int argc, char * argv[])
{
  rclcpp::init(argc, argv);
  rclcpp::spin(std::make_shared<FrameListener>());
  rclcpp::shutdown();
  return 0;
}
